import os
import json
import numpy as np


def delete_dir_if_exist(dire):
    if os.path.isdir(dire):
        command = 'rm -rf %s' % dire
        print(command)
        os.system(command)


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def get_recall(dataset_name, efsearch_l, estimate_l, time_l):
    gnd, dim = ivecs_read("../data/%s/gnd.ivecs" % dataset_name)
    topk = len(gnd[0])
    efsearch_recall_l = []
    for i, ef in enumerate(efsearch_l, 0):
        estimate = estimate_l[i]
        tmp_recall_l = []
        # for j in range(len(gnd)):
        for j in range(200):
            matches = len(np.intersect1d(gnd[j], estimate[j]))
            print(dataset_name, len(gnd), len(estimate))
            recall = float(matches / topk)
            tmp_recall_l.append(recall)
        efsearch_recall_l.append({
            "time": time_l[i],
            "recall": np.mean(tmp_recall_l),
            "efsearch": ef
        })
    return efsearch_recall_l


with open('query_log.txt', 'r') as f:
    text = f.read()
    block = text.split("\n\n\n")
with open('dataset_info.json', 'r') as f:
    ds_info = json.load(f)

ef_search_l = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280, 300, 400, 500,
               600, 700, 800, 900, 2000, 4000, 6000, 8000]
# ef_search_l = [10, 20]

delete_dir_if_exist('../result-mobius')
os.mkdir('../result-mobius')

for ds_name in ds_info:
    # for ds_name in ['audio']:
    predict_l = []
    time_l = []
    for j, efsearch in enumerate(ef_search_l, 0):
        predicted = block[j].split("\n")
        # queries per second 23140.113387, average return_percent 0.454625%
        log_info = predicted[-1]
        query_per_second = float(log_info.split(',')[0].split(' ')[-1])
        time = 1000 / query_per_second
        time_l.append(time)

        predicted = predicted[1:-1]
        predicted = np.loadtxt(predicted)
        predict_l.append(predicted)
    print("%s complete" % ds_name)
    time_recall_l = get_recall(ds_name, ef_search_l, predict_l, time_l)

    with open("../result-mobius/%s.csv" % ds_name, 'w') as f:
        f.write('ef, recall, query_time\n')
        for i in range(len(time_recall_l)):
            f.write("{}, {}, {}\n".format(time_recall_l[i]['efsearch'], time_recall_l[i]['recall'],
                                          time_recall_l[i]['time']))
    print("save %s" % ds_name)
